import cv2
import os
import numpy as np
import torch
from torch.utils.data import Dataset
import scipy.io
import pickle
import glob
from torchvision import transforms

mean = [0.443, 0.39, 0.426]
std = [0.282, 0.297, 0.281]

class OxfordRetriData(Dataset):
    def __init__(self, oxford_path):
        super().__init__()
        self.images = glob.glob(os.path.join(oxford_path, "jpg/*.jpg"))
        # gt is a list ['gnd', 'imlist', 'qimlist']
        # self.gt =  pickle.load(open(os.path.join(oxford_path, 'gnd_roxford5k.pkl'), 'rb'))
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])
    
    def __getitem__(self, i):
        img = cv2.imread(self.images[i])
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = self.transform(img)
        return img

    def __len__(self):
        return len(self.images)

def loadoxford(path):
    return OxfordRetriData(path)

if __name__ == "__main__":
    oxford_path = '/workspace/wzj/revisitop/data/datasets/roxford5k'
    # gt = os.path.join(oxford_path, 'gnd_roxford5k.pkl')
    # gt = open(gt, 'rb')
    # inf = pickle.load(gt)
    # a = ['gnd', 'imlist', 'qimlist']
    # print(inf[a[0]][1])
    # data = OxfordRetriData(oxford_path)
    # print(len(data))

    
    # labels = scipy.io.loadmat('/workspace/wzj/dataset/image_retrieval/Oxford_Buildings/gnd_roxford5k.mat')
    # print(labels.keys())
